iT邦幫忙

2023 iThome 鐵人賽

DAY 9
1

我們今天要來寫code測試Local的DynamoDB,我們找一個實際的例子來測試,目前想到會有需要存的就是使用者Google OAuth的token,我們不會希望每次要調用Google Drive的時候,都一直要你登入google進行授權,因此我們會把授權成功後拿到的token資訊存到DynamoDB去。

  1. 首先第一步就是開一個資料夾,建立4個檔案,dynamodb.go放一些操作dynamodb通用的一些function,oauth.go放處理剛剛說的token的相關方法,並建立對應的Unit test。

    https://ithelp.ithome.com.tw/upload/images/20230924/201159900NITwEAyJj.png

  2. 接著到oauth.go建立GoogleOAuthToken的struct,關於如何取得真正的token,我們到真的處理Google OAuth的部份在說明,今天先確保到DynamoDB沒問題就好。

    // oauth.go
    type GoogleOAuthToken struct {
    	PK           string `dynamodbav:"PK"`
    	AccessToken  string `dynamodbav:"access_token"`
    	TokenType    string `dynamodbav:"token_type"`
    	RefreshToken string `dynamodbav:"refresh_token"`
    	Expiry       string `dynamodbav:"expiry"`
    }
    
  3. 知道我們要操作的結構後,我們回到dynamodb.go,創建TableBasics最為基礎結構

    // dynamodb.go
    package dynamodb
    
    import (
      "context"
    	"errors"
    	"fmt"
    	"log"
    
    	"github.com/aws/aws-sdk-go-v2/aws"
    	"github.com/aws/aws-sdk-go-v2/config"
    	"github.com/aws/aws-sdk-go-v2/credentials"
    	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
    	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
    )
    
    type TableBasics struct {
    	DynamoDbClient *dynamodb.Client
    	TableName      string
    }
    
    func NewTableBasics(tableName string) *TableBasics {
    	cfg, err := config.LoadDefaultConfig(context.TODO())
    	if err != nil {
    		panic(err)
    	}
    
    	client := dynamodb.NewFromConfig(cfg)
    	return &TableBasics{
    		DynamoDbClient: client,
    		TableName:      tableName,
    	}
    }
    
  4. 下面加上CreateLocalClient,建立一個連線到地端的*dynamodb.Client,如果是平時在本機端開發,我們使用這個function得到的Client

    // dynamodb.go
    func CreateLocalClient(port int) *dynamodb.Client {
    	cfg, err := config.LoadDefaultConfig(context.TODO(),
    		config.WithRegion("ap-northeast-1"),
    		config.WithCredentialsProvider(credentials.StaticCredentialsProvider{
    			Value: aws.Credentials{
    				AccessKeyID: "dummy", SecretAccessKey: "dummy", SessionToken: "dummy",
    				Source: "Hard-coded credentials; values are irrelevant for local DynamoDB",
    			},
    		}),
    	)
    	if err != nil {
    		panic(err)
    	}
    	dsn := fmt.Sprintf("http://localhost:%d/", port)
    	return dynamodb.NewFromConfig(cfg, func(o *dynamodb.Options) {
    		o.BaseEndpoint = aws.String(dsn)
    	})
    }
    
  5. 接著打開dynamodb_test.go,我們將其作為測試的主要進入口(TestMain),建立一個TableBasics結構,其中TableName存"google-oauth”,再把原本的DynamoDbClient替換成CreateLocalClient產生的,讓他連線到我們昨天開在8000port的local dynamodb,這樣之後其他unit test直接操作初始化好的testTableBasics就好了~

    // dynamodb_test.go
    package dynamodb
    
    import (
    	"os"
    	"testing"
    )
    
    var testTableBasics *TableBasics
    
    func TestMain(m *testing.M) {
    	// google_oauth
    	testTableBasics = NewTableBasics("google-oauth")
    	// change to local dynamodb
    	testTableBasics.DynamoDbClient = CreateLocalClient(8000)
    	os.Exit(m.Run())
    }
    
  6. 再來回到oauth.go,補上建立table的方法 (從TableBasics拿出TableName,調用DynamoDbClient.CreateTable來建立table)

    //oauth.go
    package dynamodb
    
    import (
    	"context"
    	"log"
    	"time"
    
    	"github.com/aws/aws-sdk-go-v2/aws"
    	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/attributevalue"
    	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/expression"
    	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
    	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
    )
    
    type GoogleOAuthToken struct {
    	PK           string `dynamodbav:"PK"`
    	AccessToken  string `dynamodbav:"access_token"`
    	TokenType    string `dynamodbav:"token_type"`
    	RefreshToken string `dynamodbav:"refresh_token"`
    	Expiry       string `dynamodbav:"expiry"`
    }
    
    func (basics TableBasics) CreateGoogleOAuthTable() (*types.TableDescription, error) {
    	var tableDesc *types.TableDescription
    	table, err := basics.DynamoDbClient.CreateTable(context.TODO(), &dynamodb.CreateTableInput{
    
    		AttributeDefinitions: []types.AttributeDefinition{
    			{
    				AttributeName: aws.String("PK"),
    				AttributeType: types.ScalarAttributeTypeS,
    			},
    		},
    		KeySchema: []types.KeySchemaElement{{
    			AttributeName: aws.String("PK"),
    			KeyType:       types.KeyTypeHash,
    		}},
    		TableName: aws.String(basics.TableName),
    		ProvisionedThroughput: &types.ProvisionedThroughput{
    			ReadCapacityUnits:  aws.Int64(10),
    			WriteCapacityUnits: aws.Int64(10),
    		},
    	})
    	if err != nil {
    		log.Printf("Couldn't create table %v. Here's why: %v\n", basics.TableName, err)
    	} else {
    		waiter := dynamodb.NewTableExistsWaiter(basics.DynamoDbClient)
    		err = waiter.Wait(context.TODO(), &dynamodb.DescribeTableInput{
    			TableName: aws.String(basics.TableName)}, 5*time.Minute)
    		if err != nil {
    			log.Printf("Wait for table exists failed. Here's why: %v\n", err)
    		}
    		tableDesc = table.TableDescription
    	}
    	return tableDesc, err
    }
    
  7. 接著我們將PK設計為Line的User UUID,因此我們會在碰到PK的地方都加上"LINEID#”的prefix,然後補上對應的CRUD就好哩

    //oauth.go
    const PK_PREFIX_LINE = "LINEID#"
    
    // Get PK Key (with Line prefix)
    func (tok GoogleOAuthToken) GetKey() map[string]types.AttributeValue {
    	// Add prefix to PK
    	line_userid, err := attributevalue.Marshal(PK_PREFIX_LINE + tok.PK)
    	if err != nil {
    		panic(err)
    	}
    
    	return map[string]types.AttributeValue{"PK": line_userid}
    }
    
  8. 新增GoogleOAuthToken

    //oauth.go
    func (basics TableBasics) AddGoogleOAuthToken(tok GoogleOAuthToken) error {
    	tok.PK = PK_PREFIX_LINE + tok.PK
    	item, err := attributevalue.MarshalMap(tok)
    	if err != nil {
    		panic(err)
    	}
    	_, err = basics.DynamoDbClient.PutItem(context.TODO(), &dynamodb.PutItemInput{
    		TableName: aws.String(basics.TableName), Item: item,
    	})
    	if err != nil {
    		log.Printf("Couldn't add item to table. Here's why: %v\n", err)
    	}
    	return err
    }
    
  9. 更新GoogleOAuthToken

    //oauth.go
    func (basics TableBasics) TxUpdateGoogleOAuthToken(tok GoogleOAuthToken) (*dynamodb.TransactWriteItemsOutput, error) {
    	var err error
    	var response *dynamodb.TransactWriteItemsOutput
    
    	update := expression.Set(expression.Name("refresh_token"), expression.Value(tok.RefreshToken))
    	update.Set(expression.Name("access_token"), expression.Value(tok.AccessToken))
    	expr, err := expression.NewBuilder().WithUpdate(update).Build()
    	if err != nil {
    		log.Printf("Couldn't build expression for update. Here's why: %v\n", err)
    	} else {
    		twii := &dynamodb.TransactWriteItemsInput{
    			TransactItems: []types.TransactWriteItem{
    				{
    					Update: &types.Update{
    						Key:                       tok.GetKey(),
    						TableName:                 aws.String(basics.TableName),
    						ExpressionAttributeNames:  expr.Names(),
    						ExpressionAttributeValues: expr.Values(),
    						UpdateExpression:          expr.Update(),
    					},
    				},
    			},
    		}
    		response, err = basics.DynamoDbClient.TransactWriteItems(context.TODO(), twii)
    		if err != nil {
    			log.Printf("Couldn't trasnaciton update tok %v. Here's why: %v\n", tok.PK, err)
    		}
    	}
    
    	return response, err
    }
    
  10. 取得GoogleOAuthToken

    //oauth.go
    func (basics TableBasics) GetGoogleOAuthToken(line_userid string) (GoogleOAuthToken, error) {
    	tok := GoogleOAuthToken{PK: line_userid}
    	response, err := basics.DynamoDbClient.GetItem(context.TODO(), &dynamodb.GetItemInput{
    		Key: tok.GetKey(), TableName: aws.String(basics.TableName),
    	})
    	if err != nil {
    		log.Printf("Couldn't get info about %v. Here's why: %v\n", line_userid, err)
    	} else {
    		err = attributevalue.UnmarshalMap(response.Item, &tok)
    		if err != nil {
    			log.Printf("Couldn't unmarshal response. Here's why: %v\n", err)
    		}
    	}
    	return tok, err
    }
    

到這邊基本的操作就寫好了,我們到oauth_test.go來測試一下,執行單個個別的測試,可以配合昨天http://localhost:8001/ 上的GUI看一下效果,確認一下都有正確的執行就沒問題摟~

// oauth_test.go
package dynamodb

import (
	"testing"
	"github.com/stretchr/testify/assert"
)

func TestCreateGoogleOAuthTable(t *testing.T) {
	tableDesc, err := testTableBasics.CreateGoogleOAuthTable()
	t.Log("tableDesc:", tableDesc)
	t.Log("ERROR:", err)
	assert.NoError(t, err, "Expected no error creating table")
	assert.NotNil(t, tableDesc, "Table description should not be nil")
}

func TestAddGoogleOAuthToken(t *testing.T) {
	tok := GoogleOAuthToken{
		PK:           "test1234",
		AccessToken:  "test123",
		TokenType:    "Bearer",
		RefreshToken: "test123",
		Expiry:       "2023-09-24T11:31:54.2936004+08:00",
	}
	err := testTableBasics.AddGoogleOAuthToken(tok)
	if err != nil {
		t.Log("ERROR:", err)
	}
}

func TestTxUpdateGoogleOAuthToken(t *testing.T) {
	tok := GoogleOAuthToken{
		PK:          "test1234",
		AccessToken: "test123456",
		RefreshToken: "test123456",
	
	}
	output, err := testTableBasics.TxUpdateGoogleOAuthToken(tok)
	t.Log("output:", output)
	if err != nil {
		t.Log("ERROR:", err)
	}
}

func TestGetGoogleOAuthToken(t *testing.T) {
	tok, err := testTableBasics.GetGoogleOAuthToken("test1234")
	if err != nil {
		t.Log(err)
	}
	t.Log("Get Token:", tok)
}

上一篇
Day08 串接DynamoDB-01 (Cloud+Local)
下一篇
Day10 用Go操作AWS SSM
系列文
Golang LineBot X GoogleDrive:LINE有各種限制!? 那就丟上Drive吧!30
圖片
  直播研討會
圖片
{{ item.channelVendor }} {{ item.webinarstarted }} |
{{ formatDate(item.duration) }}
直播中

尚未有邦友留言

立即登入留言